Add StreamIndex Kernel for CSA [Deepseek v4] - #5079
Conversation
…_map support - Update default block_w from 128 to 512 for optimal TPU v5p memory bandwidth and pipeline overlap. - Add stop_gradient guards on inputs and outputs in csa_streamindex_score. - Wrap kernel dispatch inside jax.shard_map in DeepseekV4Indexer for distributed multi-device SPMD training. - Update unit tests with TPU smoke test and parity tolerances across CPU interpreter and TPU hardware.
… RoPE - Chunk head accumulation in VMEM (head_chunk=32, block_w=1024) to eliminate large [Bq, H, Bw] intermediate buffers and prevent VMEM OOMs. - Target TPU systolic MXU arrays with native bfloat16 einsums and float32 accumulation. - Apply rotary embeddings directly on sequence-major [B, S, H, D] tensors, eliminating 4 HBM transpositions per layer. - Add batch-divisibility guard for shard_map SPMD dispatch.
- Remove stop_gradient barriers on csa_streamindex_score inputs and outputs. - Register jax.custom_vjp on csa_streamindex_score with nondiff_argnums for tile/scale parameters. - Forward pass executes fused Pallas TPU kernel in VMEM, storing only primal inputs (zero intermediate 4D tensor in HBM). - Backward pass evaluates reference autograd, guaranteeing exact numerical parity. - Add unit tests verifying backward gradient parity and TPU hardware backward compilation.
… kernel and tests
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a fused Pallas TPU kernel for DeepSeek-V4 CSA StreamIndex score computation, integrating it into the DeepseekV4Indexer layer with a configurable flag, and adds corresponding unit and integration tests. Feedback on these changes highlights several critical improvements: correcting the batch dimension block size in the Pallas BlockSpec to prevent out-of-bounds errors, preserving NdInitializer instances during initialization wrapping to avoid sharding mismatches, removing the "context" axis from batch sharding to prevent unnecessary All-to-All resharding overhead, and planning a fused backward Pallas kernel to avoid materializing large intermediate tensors in HBM during training.
| in_specs = [ | ||
| pl.BlockSpec((None, num_heads, block_q, head_dim), lambda b, i, j: (b, 0, i, 0)), | ||
| pl.BlockSpec((None, block_w, head_dim), lambda b, i, j: (b, j, 0)), | ||
| pl.BlockSpec((None, block_q, num_heads), lambda b, i, j: (b, i, 0)), | ||
| ] | ||
| out_specs = pl.BlockSpec((None, block_q, block_w), lambda b, i, j: (b, i, j)) |
There was a problem hiding this comment.
In Pallas, using None in the block shape of a BlockSpec indicates that the block size is equal to the full dimension size (i.e., batch_size). However, the indexing function maps the batch dimension to b (the program ID for the batch dimension, which ranges from 0 to batch_size - 1). For any batch_size > 1 and b > 0, this will attempt to slice b : b + batch_size, resulting in an out-of-bounds error during compilation or execution on TPU. To correctly partition the batch dimension across programs, the block size should be set to 1 instead of None.
in_specs = [\n pl.BlockSpec((1, num_heads, block_q, head_dim), lambda b, i, j: (b, 0, i, 0)),\n pl.BlockSpec((1, block_w, head_dim), lambda b, i, j: (b, j, 0)),\n pl.BlockSpec((1, block_q, num_heads), lambda b, i, j: (b, i, 0)),\n ]\n out_specs = pl.BlockSpec((1, block_q, block_w), lambda b, i, j: (b, i, j))| def _as_nd_init(init_fn: Any) -> Any: | ||
| """Adapts a 2/3-arg Flax initializer to a 5-arg NdInitializer if needed.""" | ||
| def wrapped(key, shape, dtype, *args, **kwargs): | ||
| try: | ||
| return init_fn(key, shape, dtype, *args, **kwargs) | ||
| except TypeError: | ||
| return init_fn(key, shape, dtype) | ||
| return wrapped |
There was a problem hiding this comment.
Unconditionally wrapping kernel_init in a plain Python function closure wrapped breaks the isinstance(kernel_init, NdInitializer) checks used in DenseGeneral and other layers in MaxText. This prevents the layers from performing logically partitioned weight initialization, which can lead to incorrect weight values or sharding mismatches. We should check if init_fn is already an instance of NdInitializer and return it directly if so.
def _as_nd_init(init_fn: Any) -> Any:\n \"\"\"Adapts a 2/3-arg Flax initializer to a 5-arg NdInitializer if needed.\"\"\"\n if isinstance(init_fn, NdInitializer):\n return init_fn\n def wrapped(key, shape, dtype, *args, **kwargs):\n try:\n return init_fn(key, shape, dtype, *args, **kwargs)\n except TypeError:\n return init_fn(key, shape, dtype)\n return wrapped| total_batch_shards = 1 | ||
| if mesh is not None: | ||
| for axis_name in ("data", "fsdp", "fsdp_transpose", "expert", "context"): | ||
| if axis_name in mesh.shape: | ||
| total_batch_shards *= mesh.shape[axis_name] | ||
| if mesh is not None and total_batch_shards > 1 and (batch_size % total_batch_shards == 0): | ||
| q_pspec = jax.sharding.PartitionSpec( | ||
| ("data", "fsdp", "fsdp_transpose", "expert", "context"), | ||
| None, | ||
| None, | ||
| None, | ||
| ) | ||
| out_pspec = jax.sharding.PartitionSpec( | ||
| ("data", "fsdp", "fsdp_transpose", "expert", "context"), | ||
| None, | ||
| None, | ||
| ) |
There was a problem hiding this comment.
The \"context\" axis is reserved for sequence/context parallelism and is not used to shard the batch dimension globally in MaxText. Including \"context\" in the batch sharding axes for shard_map will cause JAX to perform an expensive All-to-All resharding of q, compressed, and weights across the context axis when context_parallelism > 1. Removing \"context\" from the batch sharding axes avoids this unnecessary overhead.
total_batch_shards = 1\n if mesh is not None:\n for axis_name in (\"data\", \"fsdp\", \"fsdp_transpose\", \"expert\"):\n if axis_name in mesh.shape:\n total_batch_shards *= mesh.shape[axis_name]\n if mesh is not None and total_batch_shards > 1 and (batch_size % total_batch_shards == 0):\n q_pspec = jax.sharding.PartitionSpec(\n (\"data\", \"fsdp\", \"fsdp_transpose\", \"expert\"),\n None,\n None,\n None,\n )\n out_pspec = jax.sharding.PartitionSpec(\n (\"data\", \"fsdp\", \"fsdp_transpose\", \"expert\"),\n None,\n None,\n )| def _csa_streamindex_score_head_major_bwd( | ||
| softmax_scale: float, | ||
| compress_rate: int, | ||
| block_q: int | None, | ||
| block_w: int | None, | ||
| interpret: bool, | ||
| res: tuple[jax.Array, jax.Array, jax.Array], | ||
| g: jax.Array, | ||
| ) -> tuple[jax.Array, jax.Array, jax.Array]: | ||
| del block_q, block_w, interpret | ||
| q, compressed, weights = res | ||
| _, vjp_fn = jax.vjp( | ||
| functools.partial( | ||
| reference_csa_streamindex_score_head_major, | ||
| softmax_scale=softmax_scale, | ||
| compress_rate=compress_rate, | ||
| ), | ||
| q, | ||
| compressed, | ||
| weights, | ||
| ) | ||
| dq, dk, dw = vjp_fn(g) | ||
| return dq, dk, dw |
There was a problem hiding this comment.
The backward pass _csa_streamindex_score_head_major_bwd currently falls back to JAX's automatic differentiation (jax.vjp) on the reference implementation reference_csa_streamindex_score_head_major. This materializes the large intermediate [B, H, S, W] tensor in HBM during training, which defeats the memory fusion benefits of the Pallas kernel. While a custom backward Pallas kernel is complex, please consider adding a TODO or planning to implement a fused backward Pallas kernel in the future to achieve full memory savings during training.
1 participant
Description
Adds a fused Pallas TPU kernel (
csa_streamindex_score) for DeepSeek-V4 Compressed Sparse Attention (CSA) indexer scoring.Summary
@jax.custom_vjp.use_csa_streamindex_kernel: bool = Falseinbase.ymlwith automatic fallback to standard einsum.Tests
tests/unit/csa_streamindex_test.py:pallas_callvsdot_general).Reproduce:
Checklist
gemini-reviewlabel.